#include "ImuTypes.h"

#include <iostream>

namespace msf
{

namespace IMU
{

const float eps = 1e-4;

/** 
 * @brief 强制让R变成一个正交矩阵
 * @param R 待优化的旋转矩阵
 * @return 优化后的矩阵
 */
Eigen::Matrix3f NormalizeRotation(const Eigen::Matrix3f &R)
{
    // 这里关注一下
    // 1. 对于行列数一样的矩阵，Eigen::ComputeThinU | Eigen::ComputeThinV    与    Eigen::ComputeFullU | Eigen::ComputeFullV 一样
    // 2. 对于行列数不同的矩阵，例如3*4 或者 4*3 矩阵只有3个奇异向量，计算的时候如果是Thin 那么得出的UV矩阵列数只能是3，如果是full那么就是4
    // 3. thin会损失一部分数据，但是会加快计算，对于大型矩阵解算方程时，可以用thin加速得到结果
    Eigen::JacobiSVD<Eigen::Matrix3f> svd(R, Eigen::ComputeFullU | Eigen::ComputeFullV);
    return svd.matrixU() * svd.matrixV().transpose();
}



/**
 * @brief                  计算旋转角度积分量
 * 
 * @param[in] angVel       陀螺仪数据
 * @param[in] imuBias      陀螺仪偏置
 * @param[in] time         两帧间的时间差
 */
IntegratedRotation::IntegratedRotation(const Eigen::Vector3f &angVel, const Bias &imuBias, const float &time)
{
    // 得到考虑偏置后的角度旋转
    const float x = (angVel(0) - imuBias.bwx) * time;
    const float y = (angVel(1) - imuBias.bwy) * time;
    const float z = (angVel(2) - imuBias.bwz) * time;

    // 计算旋转矩阵的模值，后面用罗德里格公式计算旋转矩阵时会用到
    const float d2 = x * x + y * y + z * z;
    const float d = sqrt(d2);

    Eigen::Vector3f v;
    v << x, y, z;

    // 角度转成叉积的矩阵形式
    Eigen::Matrix3f W = Sophus::SO3f::hat(v);
    // eps = 1e-4 是一个小量，根据罗德里格斯公式求极限，后面的高阶小量忽略掉得到此式
    if (d < eps)
    {
        deltaR = Eigen::Matrix3f::Identity() + W;
        rightJ = Eigen::Matrix3f::Identity();
    }
    else
    {
        deltaR = Eigen::Matrix3f::Identity() + W * sin(d) / d + W * W * (1.0f - cos(d)) / d2;
        rightJ = Eigen::Matrix3f::Identity() - W * (1.0f - cos(d)) / d2 + W * W * (d - sin(d)) / (d2 * d);
    }
}

/** 
 * @brief 预积分类构造函数，根据输入的偏置初始化预积分参数
 * @param b_ 偏置
 * @param calib imu标定参数的类
 */
Preintegrated::Preintegrated(const Bias &b_, const Calib &calib)
{
    Nga = calib.Cov;
    NgaWalk = calib.CovWalk;
    Initialize(b_);
}

/** 
 * @brief 初始化预积分
 * @param b_ 偏置
 */
void Preintegrated::Initialize(const Bias &b_)
{
    dR.setIdentity();
    dV.setZero();
    dP.setZero();
    JRg.setZero();
    JVg.setZero();
    JVa.setZero();
    JPg.setZero();
    JPa.setZero();
    C.setZero();
    Info.setZero();
    db.setZero();
    b = b_;
    bu = b_;  // 更新后的偏置
    avgA.setZero();  // 平均加速度
    avgW.setZero();  // 平均角速度
    dT = 0.0f;
    mvMeasurements.clear();  // 存放imu数据及dt
}

/** 
 * @brief 根据新的偏置重新积分mvMeasurements里的数据 Optimizer::InertialOptimization调用
 */ 
void Preintegrated::Reintegrate()
{
    std::unique_lock<std::mutex> lock(mMutex);
    const std::vector<integrable> aux = mvMeasurements;
    Initialize(bu);
    for (size_t i = 0; i < aux.size(); i++)
        IntegrateNewMeasurement(aux[i].a, aux[i].w, aux[i].t);
}

/**
 * @brief 预积分计算，更新noise
 * 
 * @param[in] acceleration  加速度计数据
 * @param[in] angVel        陀螺仪数据
 * @param[in] dt            两图像 帧之间时间差
 */
void Preintegrated::IntegrateNewMeasurement(const Eigen::Vector3f &acceleration, const Eigen::Vector3f &angVel, const float &dt)
{
    // 保存imu数据，利用中值积分的结果构造一个预积分类保存在mvMeasurements中
    mvMeasurements.push_back(integrable(acceleration, angVel, dt));

    // Position is updated firstly, as it depends on previously computed velocity and rotation.
    // Velocity is updated secondly, as it depends on previously computed rotation.
    // Rotation is the last to be updated.

    // Matrices to compute covariance
    // Step 1.构造协方差矩阵
    // 噪声矩阵的传递矩阵，这部分用于计算i到j-1历史噪声或者协方差
    Eigen::Matrix<float, 9, 9> A;
    A.setIdentity();
    // 噪声矩阵的传递矩阵，这部分用于计算j-1新的噪声或协方差，这两个矩阵里面的数都是当前时刻的，计算主要是为了下一时刻使用
    Eigen::Matrix<float, 9, 6> B;
    B.setZero();

    // 考虑偏置后的加速度、角速度
    Eigen::Vector3f acc, accW;
    acc << acceleration(0) - b.bax, acceleration(1) - b.bay, acceleration(2) - b.baz;
    accW << angVel(0) - b.bwx, angVel(1) - b.bwy, angVel(2) - b.bwz;

    // 记录平均加速度和角速度
    avgA = (dT * avgA + dR * acc * dt) / (dT + dt);
    avgW = (dT * avgW + accW * dt) / (dT + dt);

    // Update delta position dP and velocity dV (rely on no-updated delta rotation)
    // 根据没有更新的dR来更新dP与dV  eq.(38)
    dP = dP + dV * dt + 0.5f * dR * acc * dt * dt;
    dV = dV + dR * acc * dt;

    // Compute velocity and position parts of matrices A and B (rely on non-updated delta rotation)
    // 根据η_ij = A * η_i,j-1 + B_j-1 * η_j-1中的Ａ矩阵和Ｂ矩阵对速度和位移进行更新
    Eigen::Matrix<float, 3, 3> Wacc = Sophus::SO3f::hat(acc);

    A.block<3, 3>(3, 0) = -dR * dt * Wacc;
    A.block<3, 3>(6, 0) = -0.5f * dR * dt * dt * Wacc;
    A.block<3, 3>(6, 3) = Eigen::DiagonalMatrix<float, 3>(dt, dt, dt);
    B.block<3, 3>(3, 3) = dR * dt;
    B.block<3, 3>(6, 3) = 0.5f * dR * dt * dt;

    // Update position and velocity jacobians wrt bias correction
    // 因为随着时间推移，不可能每次都重新计算雅克比矩阵，所以需要做J(k+1) = j(k) + (~)这类事，分解方式与AB矩阵相同
    // 论文作者对forster论文公式的基础上做了变形，然后递归更新，参见 https://github.com/UZ-SLAMLab/ORB_SLAM3/issues/212
    JPa = JPa + JVa * dt - 0.5f * dR * dt * dt;
    JPg = JPg + JVg * dt - 0.5f * dR * dt * dt * Wacc * JRg;
    JVa = JVa - dR * dt;
    JVg = JVg - dR * dt * Wacc * JRg;

    // Update delta rotation
    // Step 2. 构造函数，会根据更新后的bias进行角度积分
    IntegratedRotation dRi(angVel, b, dt);
    // 强行归一化使其符合旋转矩阵的格式
    dR = NormalizeRotation(dR * dRi.deltaR);

    // Compute rotation parts of matrices A and B
    // 补充AB矩阵
    A.block<3, 3>(0, 0) = dRi.deltaR.transpose();
    B.block<3, 3>(0, 0) = dRi.rightJ * dt;

    // 小量delta初始为0，更新后通常也为0，故省略了小量的更新
    // Update covariance
    // Step 3.更新协方差，frost经典预积分论文的第63个公式，推导了噪声（ηa, ηg）对dR dV dP 的影响
    C.block<9, 9>(0, 0) = A * C.block<9, 9>(0, 0) * A.transpose() + B * Nga * B.transpose();  // B矩阵为9*6矩阵 Nga 6*6对角矩阵，3个陀螺仪噪声的平方，3个加速度计噪声的平方
    // 这一部分最开始是0矩阵，随着积分次数增加，每次都加上随机游走，偏置的信息矩阵
    C.block<6, 6>(9, 9) += NgaWalk;

    // Update rotation jacobian wrt bias correction
    // 计算偏置的雅克比矩阵，r对bg的导数，∂ΔRij/∂bg = (ΔRjj-1) * ∂ΔRij-1/∂bg - Jr(j-1)*t
    // 论文作者对forster论文公式的基础上做了变形，然后递归更新，参见 https://github.com/UZ-SLAMLab/ORB_SLAM3/issues/212
    // ? 为什么先更新JPa、JPg、JVa、JVg最后更新JRg? 答：这里必须先更新dRi才能更新到这个值，但是为什么JPg和JVg依赖的上一个JRg值进行更新的？
    JRg = dRi.deltaR.transpose() * JRg - dRi.rightJ * dt;

    // Total integrated time
    // 更新总时间
    dT += dt;
}

/** 
 * @brief 更新偏置
 * @param bu_ 偏置
 */
void Preintegrated::SetNewBias(const Bias &bu_)
{
    std::unique_lock<std::mutex> lock(mMutex);
    bu = bu_;

    db(0) = bu_.bwx - b.bwx;
    db(1) = bu_.bwy - b.bwy;
    db(2) = bu_.bwz - b.bwz;
    db(3) = bu_.bax - b.bax;
    db(4) = bu_.bay - b.bay;
    db(5) = bu_.baz - b.baz;
}

/** 
 * @brief 获得当前偏置与输入偏置的改变量
 * @param b_ 偏置
 * @return 改变量
 */
IMU::Bias Preintegrated::GetDeltaBias(const Bias &b_)
{
    std::unique_lock<std::mutex> lock(mMutex);
    return IMU::Bias(b_.bax - b.bax, b_.bay - b.bay, b_.baz - b.baz, b_.bwx - b.bwx, b_.bwy - b.bwy, b_.bwz - b.bwz);
}

/** 
 * @brief 根据新的偏置计算新的dR
 * @param b_ 新的偏置
 * @return dR
 */
Eigen::Matrix3f Preintegrated::GetDeltaRotation(const Bias &b_)
{
    std::unique_lock<std::mutex> lock(mMutex);
    // 计算偏置的变化量
    Eigen::Vector3f dbg;

    dbg << b_.bwx - b.bwx, b_.bwy - b.bwy, b_.bwz - b.bwz;

    // 考虑偏置后，dR对偏置线性化的近似求解,邱笑晨《预积分总结与公式推导》P13～P14
    // Forster论文公式（44）yP17也有结果（但没有推导），后面两个函数GetDeltaPosition和GetDeltaPosition也是基于此推导的
    return NormalizeRotation(dR * Sophus::SO3f::exp(JRg * dbg).matrix());
}

/** 
 * @brief 根据新的偏置计算新的dV
 * @param b_ 新的偏置
 * @return dV
 */
Eigen::Vector3f Preintegrated::GetDeltaVelocity(const Bias &b_)
{
    std::unique_lock<std::mutex> lock(mMutex);
    Eigen::Vector3f dbg, dba;
    dbg << b_.bwx - b.bwx, b_.bwy - b.bwy, b_.bwz - b.bwz;
    dba << b_.bax - b.bax, b_.bay - b.bay, b_.baz - b.baz;
    // 考虑偏置后，dV对偏置线性化的近似求解,邱笑晨《预积分总结与公式推导》P13，JPg和JPa在预积分处理中更新 
    return dV + JVg * dbg + JVa * dba;
}

/** 
 * @brief 根据新的偏置计算新的dP
 * @param b_ 新的偏置
 * @return dP
 */
Eigen::Vector3f Preintegrated::GetDeltaPosition(const Bias &b_)
{
    std::unique_lock<std::mutex> lock(mMutex);
    Eigen::Vector3f dbg, dba;
    dbg << b_.bwx - b.bwx, b_.bwy - b.bwy, b_.bwz - b.bwz;
    dba << b_.bax - b.bax, b_.bay - b.bay, b_.baz - b.baz;
    // 考虑偏置后，dP对偏置线性化的近似求解,邱笑晨《预积分总结与公式推导》P13，JPg和JPa在预积分处理中更新
    return dP + JPg * dbg + JPa * dba;
}

/** 
 * @brief 返回经过db(δba, δbg)更新后的dR,与上面是一个意思
 * @return dR
 */
Eigen::Matrix3f Preintegrated::GetUpdatedDeltaRotation()
{
    std::unique_lock<std::mutex> lock(mMutex);
    return NormalizeRotation(dR * Sophus::SO3f::exp(JRg * db.head(3)).matrix());
}

/** 
 * @brief 返回经过db(δba, δbg)更新后的dV,与上面是一个意思
 * @return dV
 */
Eigen::Vector3f Preintegrated::GetUpdatedDeltaVelocity()
{
    std::unique_lock<std::mutex> lock(mMutex);
    return dV + JVg * db.head(3) + JVa * db.tail(3);
}

/** 
 * @brief 返回经过db(δba, δbg)更新后的dP,与上面是一个意思
 * @return dP
 */
Eigen::Vector3f Preintegrated::GetUpdatedDeltaPosition()
{
    std::unique_lock<std::mutex> lock(mMutex);
    return dP + JPg * db.head(3) + JPa * db.tail(3);
}

/** 
 * @brief 设置参数
 * @param Tbo_ 位姿变换
 * @param ng 噪声
 * @param na 噪声
 * @param ngw 随机游走
 * @param naw 随机游走
 */
void Calib::Set(const Sophus::SE3<float> &sophTbo, const float &ng, const float &na, const float &ngw, const float &naw)
{
    mbIsSet = true;
    const float ng2 = ng * ng;
    const float na2 = na * na;
    const float ngw2 = ngw * ngw;
    const float naw2 = naw * naw;

    // Sophus/Eigen
    mTbo = sophTbo;
    mTob = mTbo.inverse();
    // 噪声协方差
    Cov.diagonal() << ng2, ng2, ng2, na2, na2, na2;
    // 随机游走协方差
    CovWalk.diagonal() << ngw2, ngw2, ngw2, naw2, naw2, naw2;
}

/** 
 * @brief imu标定参数的构造函数
 * @param calib imu标定参数
 */
Calib::Calib(const Calib &calib)
{
    mbIsSet = calib.mbIsSet;
    // Sophus/Eigen parameters
    mTbo = calib.mTbo;
    mTob = calib.mTob;
    Cov = calib.Cov;
    CovWalk = calib.CovWalk;
}

// ESKF
void propagate_state(
        IMU::Point last_imu,
        IMU::Point curr_imu,
        const State &last_state,
        State &state,
        const Calib &calib,
        bool with_noise,
        bool with_const_noise,
        const Eigen::Matrix<double, kNoiseDim, 1> &vec_noise) 
{
    const double dt = curr_imu.t - last_imu.t;
    const double dt2 = dt * dt;

    Eigen::Vector3d vec_na = Eigen::Vector3d::Zero();
    Eigen::Vector3d vec_ng = Eigen::Vector3d::Zero();
    Eigen::Vector3d vec_wa = Eigen::Vector3d::Zero();
    Eigen::Vector3d vec_wg = Eigen::Vector3d::Zero();
    if (with_noise) {
        // TODO: check vec_noise empty or not
        if (with_const_noise) {
            vec_na = Eigen::Vector3d(calib.Cov.diagonal()[3], calib.Cov.diagonal()[3], calib.Cov.diagonal()[3]); // 加速度计随机(测量)噪声
            vec_ng = Eigen::Vector3d(calib.Cov.diagonal()[0], calib.Cov.diagonal()[0], calib.Cov.diagonal()[0]); // 陀螺仪随机(测量)噪声
            vec_wa = Eigen::Vector3d(calib.CovWalk.diagonal()[3], calib.CovWalk.diagonal()[3], calib.CovWalk.diagonal()[3]); // 加速度计漂移(零偏)噪声
            vec_wg = Eigen::Vector3d(calib.CovWalk.diagonal()[0], calib.CovWalk.diagonal()[0], calib.CovWalk.diagonal()[0]); // 陀螺仪漂移(零偏)噪声
        } else {
            vec_na = vec_noise.segment<3>(0);
            vec_ng = vec_noise.segment<3>(3);
            vec_wa = vec_noise.segment<3>(6);
            vec_wg = vec_noise.segment<3>(9);
        }
    }

    const Eigen::Vector3d acc_unbias = 0.5 * (last_imu.a.cast<double>() + curr_imu.a.cast<double>()) - last_state.acc_bias - vec_na;
    const Eigen::Vector3d gyr_unbias = 0.5 * (last_imu.w.cast<double>() + curr_imu.w.cast<double>()) - last_state.gyr_bias - vec_ng;

    Eigen::Vector3d acc_nominal = last_state.Rwb_ * acc_unbias + Eigen::Vector3d(0, 0, -kG);

    // Problem6 重力加速度导致漂移(直接置0，这样只能二维定位了)
    // acc_nominal[2] = 0.0;

    // 加速度约束：减去重力影响后，再进行加速度约束，最后转到世界坐标表示
    Eigen::Vector3d acc_bodyframe = last_state.Rwb_.inverse() * acc_nominal;
    //acc_bodyframe[1] = 0.0; //imu坐标系下，y方向加速度为0
    acc_bodyframe[2] = 0.0; //imu坐标系下，z方向加速度为0
    acc_nominal = last_state.Rwb_ * acc_bodyframe;

    const auto &dR = State::delta_rot_mat(gyr_unbias * dt);

    state.p_wb_ = last_state.p_wb_ + last_state.v_wb_ * dt + 0.5 * acc_nominal * dt2;
    state.v_wb_ = last_state.v_wb_ + acc_nominal * dt;
    state.Rwb_ = State::rotation_update(last_state.Rwb_, dR);
    state.acc_bias = last_state.acc_bias + vec_wa * dt;
    state.gyr_bias = last_state.gyr_bias + vec_wg * dt;
}

void propagate_state_cov(const IMU::Point last_imu, const IMU::Point curr_imu, const State &last_state, State &state, const Calib &calib) {
    const double dt = curr_imu.t - last_imu.t;

    const Eigen::Vector3d acc_unbias = 0.5 * (last_imu.a.cast<double>() + curr_imu.a.cast<double>()) - last_state.acc_bias;
    const Eigen::Vector3d gyr_unbias = 0.5 * (last_imu.w.cast<double>() + curr_imu.w.cast<double>()) - last_state.gyr_bias;

    const auto &dR = State::delta_rot_mat(gyr_unbias * dt);

    // Fx
    MatrixSD Fx = MatrixSD::Identity();
    Fx.block<3, 3>(0, 3) = Eigen::Matrix3d::Identity() * dt;
    Fx.block<3, 3>(3, 6) = -state.Rwb_ * Converter::skew_matrix(acc_unbias) * dt;
    Fx.block<3, 3>(3, 9) = -state.Rwb_ * dt;
    if (State::kAngError == ANGULAR_ERROR::LOCAL_ANGULAR_ERROR) 
    {
        Fx.block<3, 3>(6, 6) = dR.transpose();
        Fx.block<3, 3>(6, 12) = -Eigen::Matrix3d::Identity() * dt;
    } 
    else 
    {
        Fx.block<3, 3>(6, 12) = -state.Rwb_ * dt;
    }

    // P: error-state covariance
    state.cov = Fx * last_state.cov * Fx.transpose() + noise_cov_discret_time(dt,calib);
}

Eigen::Matrix<double, kNoiseDim, kNoiseDim> noise_cov(double dt, const Calib &calib) {
    const double dt2 = dt * dt;
    Eigen::Matrix<double, kNoiseDim, kNoiseDim> Qi = Eigen::Matrix<double, kNoiseDim, kNoiseDim>::Zero();
    Qi.block<3, 3>(0, 0) = Eigen::Matrix3d::Identity() * dt2 * calib.Cov.diagonal()[3]; // a noise
    Qi.block<3, 3>(3, 3) = Eigen::Matrix3d::Identity() * dt2 * calib.Cov.diagonal()[0]; // w noise
    Qi.block<3, 3>(6, 6) = Eigen::Matrix3d::Identity() * dt * calib.CovWalk.diagonal()[3]; // a walk
    Qi.block<3, 3>(9, 9) = Eigen::Matrix3d::Identity() * dt * calib.CovWalk.diagonal()[0]; // w walk
    return Qi;
}

Eigen::Matrix<double, kStateDim, kStateDim> noise_cov_discret_time(double dt, const Calib &calib) {
    Eigen::Matrix<double, kStateDim, kNoiseDim> Fi = Eigen::Matrix<double, kStateDim, kNoiseDim>::Zero();
    Fi.block<12, kNoiseDim>(3, 0) = Eigen::Matrix<double, 12, kNoiseDim>::Identity();
    return Fi * noise_cov(dt, calib) * Fi.transpose();
}


} // namespace IMU

} // namespace msf
